package com.mi.system.websocket;

import cn.hutool.core.util.StrUtil;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.mi.common.core.exception.JwtSignatureVerifyException;
import com.mi.common.security.config.JwtProperties;
import com.mi.common.security.util.JwtPayload;
import com.mi.common.security.util.JwtUtil;
import com.nimbusds.jose.JOSEException;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.text.ParseException;
import java.util.Map;


/**
 * websocket握手拦截器
 *
 * @author May
 * @version 1.0
 * @date 2018/10/24 9:28
 */
@Component
@RequiredArgsConstructor
@Slf4j
public class WebSocketHandsShakeInterceptor implements HandshakeInterceptor {

    private static final String USERNAME = "username";
    private static final String AUTHORIZATION = "Authorization";

    private final ObjectMapper objectMapper;
    private final JwtProperties jwtProperties;

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,
                                   Map<String, Object> attributes) {

        log.info("Authentication Check Before WebSocket Connect's Hand Shake.");

        // http协议转换webSocket协议进行前，token作为用户凭证
        ServletServerHttpRequest httpRequest = (ServletServerHttpRequest) request;
        String authToken = httpRequest.getServletRequest().getParameter(AUTHORIZATION);

        if (StrUtil.isEmpty(authToken)) {
            return false;
        }

        try {
            String payload = JwtUtil.verifySignature(authToken, jwtProperties.getSecret());
            JwtPayload jwtPayload = objectMapper.readValue(payload, JwtPayload.class);
            if (jwtPayload.isExpired()) {
                return false;
            }
            attributes.put(USERNAME, jwtPayload.getUsername());
            request.getHeaders().add(USERNAME, jwtPayload.getUsername());
            log.info("Websocket Connect Authentication Pass : {}", jwtPayload.getUsername());
        } catch (JOSEException | ParseException | JwtSignatureVerifyException e) {
            log.error("User Token Verify Error : {}", e.getMessage());
            return false;
        } catch (Exception e) {
            log.error("UnKnown Exception : {}", e.getMessage());
            return false;
        }

        return true;
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,
                               Exception ex) {
        //握手成功后
        log.info("After WebSocket Connect Hand Shake Success Switch HTTP TO WebSocket.");
        response.setStatusCode(HttpStatus.SWITCHING_PROTOCOLS);
    }
}
